import matplotlib.pyplot as plt
import numpy as np

'''------------------------------------------- Residual-based adaptive refinement ---------------------------------------------'''
def get_top_k_indices(arr, k=100):
    '''get indices of the top k values in arr'''
    arr = arr.reshape(-1)  # reshape the array to 1D
    indices = np.argsort(arr)[::-1]  # sort indices in descending order
    return indices[:k]

def plot_room_add_points_3d(points):
    fig = plt.figure(dpi=200)
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0.5)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')  
    # plt.axis("scaled")
    plt.show()

def convert_and_remove_duplicates(array):
    # Convert array of int32 to list
    array_list = [list(arr) for arr in array]
    # Flatten the 2D list
    flattened_list = [item for sublist in array_list for item in sublist]
    # Remove duplicate elements
    unique_list = list(set(flattened_list))
    return unique_list

def plot_points_xy(points, color="k", marker="."):
    figure = plt.figure()
    axis = figure.add_subplot(111)
    axis.scatter(points[:, 0], points[:, 1], s=0.1, color=color, marker=marker)
    # plt.axis("scaled")  # 设置x轴和y轴相同的缩放比例
    plt.show()

def tetrahedron_centroid(tetrahedrons, vertices):
    centroids = []
    for tetrahedron in tetrahedrons:
        vertex1 = vertices[tetrahedron[0]]
        vertex2 = vertices[tetrahedron[1]]
        vertex3 = vertices[tetrahedron[2]]
        vertex4 = vertices[tetrahedron[3]]
        centroid = (vertex1 + vertex2 + vertex3 + vertex4) / 4
        centroids.append(centroid)
    centroids_points = np.array(centroids)    
    return centroids_points

def tetrahedron_midpoint(tetrahedrons, vertices):
    midpoints = []
    for tetrahedron in tetrahedrons:
        vertex1 = vertices[tetrahedron[0]]
        vertex2 = vertices[tetrahedron[1]]
        vertex3 = vertices[tetrahedron[2]]
        vertex4 = vertices[tetrahedron[3]]
        midpoint12 = (vertex1 + vertex2) / 2
        midpoints.append(midpoint12)
        midpoint13 = (vertex1 + vertex3) / 2
        midpoints.append(midpoint13)
        midpoint14 = (vertex1 + vertex4) / 2
        midpoints.append(midpoint14)
        midpoint23 = (vertex2 + vertex3) / 2
        midpoints.append(midpoint23)
        midpoint24 = (vertex2 + vertex4) / 2
        midpoints.append(midpoint24)
        midpoint34 = (vertex3 + vertex4) / 2
        midpoints.append(midpoint34)
    midpoints = np.array(midpoints)    
    return midpoints

def remove_duplicate_points(points):
    unique_points = np.unique(points, axis=0)
    return unique_points